#include <iostream>
#include <fstream>
#include <string>
#include <vector>

static const std::vector<std::string> g_desc_list = {
    "Component::Info",
    "Component::Write",
    "Component::CompressWrite",
    "Component::Backprop",
    "Component::ConsolidateMemory",
    "Component::Check",
    "Component::UpdateSimple",
    "Component::UpdateNaturalGradient",
    "Component::UpdateParameters",
    "Component::FreezeNaturalGradient"
};

static std::vector<std::string> g_def_list = {
    "std::string Info",
    "void Write",
    "void CompressWrite",
    "void Backprop",
    "void ConsolidateMemory",
    "void Check",
    "void UpdateSimple",
    "void UpdateNaturalGradient",
    "void UpdateParameters",
    "void FreezeNaturalGradient"
};

bool is_match(const std::string &line, const std::vector<std::string> &func_list, bool head = true, bool tail = true) {
    for (auto func : func_list) {
        size_t pos = line.find(func);
        if (pos == std::string::npos) continue;
        if (head && pos != 0 && std::isalnum(line[pos-1])) continue; // alpha or num before
        if (tail && std::isalnum(line[pos + func.length()])) continue;
        return true;
    }
    return false;
}

bool is_description(const std::string &line) {
    return is_match(line, g_desc_list, false);
}
bool is_definition(const std::string &line) {
    return is_match(line, g_def_list);
}

bool statment_end(const std::string &line, int &depth, char start, char end) {
    bool reached = (depth != 0);
    const char *s = line.c_str(), *e = s + line.length();
    while (s < e) {
        if (*s == start) {
            depth ++;
            reached = true;
        }
        if (*s == end) depth --;
        s++;
    }
    return reached && depth == 0;
}

void write_statments(const std::string &line, std::ifstream& ifs, std::ofstream &ofs, char start, char end) {
    static const std::string head = "#ifndef KALDI_LITE\n", tail = "#endif\n";
    ofs.write(head.c_str(), head.length());
    ofs.write(line.c_str(), line.length());
    int depth = 0;
    if (start != ' ' && end != ' ' && !statment_end(line, depth, start, end)) {
        std::string new_line;
        while(std::getline(ifs, new_line)) {
            new_line += "\n";
            ofs.write(new_line.c_str(), new_line.length());
            if (statment_end(new_line, depth, start, end)) break;
        }
    }
    ofs.write(tail.c_str(), tail.length());
}

bool is_variable(const std::string &line) {
    return line.find("OnlineNaturalGradient preconditioner_;") != std::string::npos ||
        line.find("OnlineNaturalGradient preconditioner_in_;") != std::string::npos ||
        line.find("OnlineNaturalGradient preconditioner_out_;") != std::string::npos ||
        line.find("OnlineNaturalGradient scale_preconditioner_;") != std::string::npos ||
        line.find("OnlineNaturalGradient offset_preconditioner_;") != std::string::npos;
}

void prune(const std::string &input) {
    const std::string output = input + "_new";
    std::ifstream ifs(input);
    std::ofstream ofs(output);
    std::string line;
    while (std::getline(ifs, line)) {
        line += "\n";
        if (line == "#ifndef KALDI_LITE\n") {
            ofs.write(line.c_str(), line.length());
            int depth = 1;
            while (std::getline(ifs, line)) {
                line += "\n";
                ofs.write(line.c_str(), line.length());
                if (line.find("#if") != std::string::npos) depth++;
                if (line.find("#endif") != std::string::npos) depth--;
                if (line == "#endif\n" && depth == 0) break;
            }
            continue;
        }
        if (is_description(line)) {
            write_statments(line, ifs, ofs, '{', '}');
        } else if (is_definition(line)) {
            write_statments(line, ifs, ofs, '(', ')');
        } else if (is_variable(line)) {
            write_statments(line, ifs, ofs, ' ', ' ');
        } else {
            ofs.write(line.c_str(), line.length());
        }
    }
    ofs.close();
    ifs.close();
    std::rename(output.c_str(), input.c_str());
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        std::cout << "nnet3-prune file" << std::endl;
        return 1;
    }
    const char *file = argv[1];
    prune(file);
    return 0;
}
